require "bit-struct"

module MediaTools

  class ChunkHeader < BitStruct
    default_options :endian => :little
    
    char :fourcc, 32, "fourCC code"
    unsigned :chunk_length, 32, "chunk length"
  end

  class RiffChunkHeader < ChunkHeader
    char :format, 32, "format (WAVE, typically)"
    
    initial_value.fourcc = "RIFF"
  end
  
  class PcmAudioFormat < BitStruct
    default_options :endian => :little
    
    unsigned :format, 16, " PCM = 1 (i.e. Linear quantization), values other than 1 indicate some form of compression."
    unsigned :num_channels, 16, " Mono = 1, Stereo = 2, etc."
    unsigned :sample_rate, 32, "8000, 44100, etc."
    unsigned :byte_rate, 32, "== SampleRate * NumChannels * BitsPerSample/8"
    unsigned :block_align, 16, " == NumChannels * BitsPerSample/8"
    unsigned :bits_per_sample, 16, "8 bits = 8, 16 bits = 16, etc."
    
    def audio_chunk_class
      unless @audio_chunk_class
        format = self
        @audio_chunk_class = Class.new(BitStruct)
        @audio_chunk_class.class_eval do
          default_options :endian => :little
          
          (0..format.num_channels - 1).each do |channel|
            send :unsigned, "channel#{channel}".to_sym, format.bits_per_sample, "channel #{channel} data"
          end
          
          define_method :channels do
             (0..format.num_channels - 1).map do |index|
               send "channel#{index}".to_sym
             end
          end
          
          self
        end
      end
      @audio_chunk_class
    end
  end
  
  module RiffReader
    def read_chunk io
      buffer = io.read(ChunkHeader.round_byte_length)
      chunk_header = ChunkHeader.new buffer
      header_pos = io.pos
      on_chunk chunk_header, io # callback
      after_pos = header_pos + chunk_header.chunk_length
      after_pos += after_pos % 2
      io.pos = after_pos
    end
    private :read_chunk
    
    def read_all io
      buffer = io.read(RiffChunkHeader.round_byte_length)
      riff_chunk_header = RiffChunkHeader.new buffer
      on_riff_chunk_header riff_chunk_header if respond_to? :on_riff_chunk_header, true # callback
      
      while !io.eof?
        read_chunk io
      end
      on_eof if respond_to? :on_eof, true # callback
    end
  end
  
  module MicrosoftWavPcmReader
    include RiffReader
    
    attr :format
    
    def on_chunk(header, io)
      case header.fourcc
      when "fmt ":
        buffer = io.read(PcmAudioFormat.round_byte_length)
        @format = PcmAudioFormat.new(buffer)
        on_format @format if respond_to? :on_format, true # callback
      when "data"
        if respond_to? :on_data, true
          on_data header.chunk_length, io
        elsif respond_to? :on_data_sample, true
          end_pos = io.pos + header.chunk_length
          while io.pos < end_pos
            buffer = io.read(@format.audio_chunk_class.round_byte_length)
            data_sample = @format.audio_chunk_class.new(buffer)
            on_data_sample data_sample # callback
          end
        end
      end
    end
  end
  
  class MicrosoftWavPcmFormatExtractor
    include MicrosoftWavPcmReader
  end
  
  class RiffWriter
    def initialize(io)
      @io = io
      @stack = []
    end
    
    attr_reader :io

    def write_chunk(header)
      write_chunk_begin header
      yield if block_given?
      write_chunk_end header
    end
    
    def write_chunk_begin(header)
      header_pos = @io.pos
      @io.write(header)
      start_pos = @io.pos

      @stack << [header, header_pos, start_pos]
    end
    
    def write_chunk_end(header)
      begin 
        sd = @stack.pop

        end_pos = @io.pos
        sd[0].chunk_length = end_pos - sd[2]
        @io.pos = sd[1]
        @io.write(sd[0])
        @io.pos = end_pos      
      end until sd.nil? || sd[0].eql?(header)
      raise "invalid header chunk specified: #{header}" unless sd
    end
  end
  
  class MicrosoftWavPcmWriter
    def initialize(io)
      @writer = RiffWriter.new io
    end
    
    def io
      @writer.io
    end

    def write_headers(format, &block)
      @writer.write_chunk(RiffChunkHeader.new(:format => "WAVE")) do
        @writer.write_chunk ChunkHeader.new(:fourcc => "fmt ") do 
          @writer.io.write(format)
        end
        @writer.write_chunk(ChunkHeader.new(:fourcc => "data"), &block)
      end
    end
    
    def write_headers_begin(format)
      @riff_chunk_header = RiffChunkHeader.new(:format => "WAVE")
      
      @writer.write_chunk_begin @riff_chunk_header
      @writer.write_chunk ChunkHeader.new(:fourcc => "fmt ") do 
        @writer.io.write(format)
      end
      @writer.write_chunk_begin ChunkHeader.new(:fourcc => "data")
    end
    
    def write_headers_end
      @writer.write_chunk_end @riff_chunk_header
    end
  end
  
end